import torch.nn as nn
import torch
import math
import numpy as np
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch_geometric.utils import erdos_renyi_graph, remove_self_loops, add_self_loops, degree, add_remaining_self_loops
from data_utils import sys_normalized_adjacency, sparse_mx_to_torch_sparse_tensor
from torch_sparse import SparseTensor, matmul


# 定义一个名为 gcn_conv 的函数，该函数实现图卷积网络（GCN）中的卷积操作。
# x：输入特征矩阵，形状为 [N, D]，其中 N 是节点数，D 是特征维度。
# edge_index：边的索引矩阵，形状为 [2, E]，其中 E 是边的数量。
def gcn_conv(x, edge_index):
    # 获取节点数 N，即 x 的第一个维度的大小。
    N = x.shape[0]
    # 从 edge_index 中提取起始节点和结束节点的索引。row 表示每条边的起始节点，col 表示每条边的结束节点。
    row, col = edge_index
    # 计算每个节点的度（degree），使用结束节点 col 和节点总数 N。
    # degree(col, N) 返回每个节点的度，类型为整型张量。
    # .float() 将其转换为浮点型张量。
    d = degree(col, N).float()
    # 计算每条边结束节点的度的倒数平方根，得到入度的归一化因子。
    # 1. / d[col] 计算结束节点的度的倒数。
    # .sqrt() 计算平方根。
    d_norm_in = (1. / d[col]).sqrt()
    # 计算每条边起始节点的度的倒数平方根，得到出度的归一化因子。
    # 1. / d[row] 计算起始节点的度的倒数。
    # .sqrt() 计算平方根。
    d_norm_out = (1. / d[row]).sqrt()
    # 创建一个与 row 形状相同的全为1的张量 value。
    # torch.ones_like(row) 返回一个与 row 形状相同的全为1的张量。
    # d_norm_in * d_norm_out 计算每条边的权重，使用起始节点和结束节点的度的倒数平方根的乘积。
    # torch.ones_like(row) * d_norm_in * d_norm_out 将计算的边的权重赋值给 value。
    value = torch.ones_like(row) * d_norm_in * d_norm_out
    # 将 value 中的 NaN、正无穷和负无穷替换为0，以防止在计算过程中出现无效值。
    # torch.nan_to_num(value, nan=0.0, posinf=0.0, neginf=0.0) 用0替换 value 中的无效值。
    value = torch.nan_to_num(value, nan=0.0, posinf=0.0, neginf=0.0)
    # 创建一个稀疏邻接矩阵 adj，其大小为 [N, N]。
    # row=col 和 col=row 交换了起始节点和结束节点的位置，这是为了构建对称的邻接矩阵。
    # value 是边的权重。
    # SparseTensor(row=col, col=row, value=value, sparse_sizes=(N, N)) 构造稀疏邻接矩阵。
    adj = SparseTensor(row=col, col=row, value=value, sparse_sizes=(N, N))
    # 执行稀疏矩阵与输入特征矩阵 x 的乘法，得到卷积后的特征矩阵。
    # matmul(adj, x) 返回形状为 [N, D] 的矩阵，表示卷积后的节点特征。
    return matmul(adj, x)  # [N, D]


# 定义一个名为 GraphConvolutionBase 的类，继承自 torch.nn.Module，这是 PyTorch 中所有神经网络模块的基类。
class GraphConvolutionBase(nn.Module):

    # 定义类的初始化方法 __init__，接受输入特征维度 in_features、输出特征维度 out_features 和一个布尔值 residual（表示是否使用残差连接）。
    # 调用父类 nn.Module 的初始化方法 super(GraphConvolutionBase, self).__init__()。
    # 将 residual、in_features 和 out_features 参数保存在实例变量中。
    def __init__(self, in_features, out_features, residual=False):
        super(GraphConvolutionBase, self).__init__()
        self.residual = residual
        self.in_features = in_features
        self.out_features = out_features
        # 定义权重矩阵 weight，形状为 [in_features, out_features]，并将其注册为模型参数。
        # 如果使用残差连接，则额外定义另一个权重矩阵 weight_r，同样注册为模型参数。
        self.weight = Parameter(torch.FloatTensor(self.in_features, self.out_features))
        if self.residual:
            self.weight_r = Parameter(torch.FloatTensor(self.in_features, self.out_features))
        # 调用 reset_parameters 方法初始化权重参数。
        self.reset_parameters()

    # 定义 reset_parameters 方法，用于初始化权重。
    # stdv 是用于初始化权重的标准差，计算方式为 1 / sqrt(out_features)。
    # 使用 uniform_(-stdv, stdv) 方法将权重初始化为均匀分布的值，范围在 [-stdv, stdv] 之间。
    # 如果使用残差连接，同样初始化 weight_r。
    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.out_features)
        self.weight.data.uniform_(-stdv, stdv)
        self.weight_r.data.uniform_(-stdv, stdv)

    # 定义前向传播方法 forward，接受输入特征矩阵 x、边索引 adj 和初始输入 x0（这里 x0 没有使用）。
    # 调用 gcn_conv(x, adj) 计算图卷积，将结果保存在 hi 中。
    # 使用 torch.mm 进行矩阵乘法，将 hi 与权重 weight 相乘，得到输出 output。
    # 如果使用残差连接，则将 x 与权重 weight_r 相乘的结果加到 output 上。
    # 返回最终的输出 output。
    def forward(self, x, adj, x0):
        hi = gcn_conv(x, adj)
        output = torch.mm(hi, self.weight)
        if self.residual:
            output = output + torch.mm(x, self.weight_r)
        return output


# 定义一个 CaNetConv 类，继承自 torch.nn.Module。
class CaNetConv(nn.Module):

    # 初始化方法 __init__ 接受输入特征维度 in_features、输出特征维度 out_features、层数 K、是否使用残差连接 residual、
    # 后端类型 backbone_type（默认为 gcn）、是否使用变体 variant、设备 device。
    def __init__(self, in_features, out_features, K, residual=True, backbone_type='gcn', variant=False, device=None):
        super(CaNetConv, self).__init__()
        self.backbone_type = backbone_type
        self.out_features = out_features
        self.residual = residual
        if backbone_type == 'gcn':
            # 如果 backbone_type 是 gcn，则初始化一个权重矩阵 self.weights，其形状为 [K, in_features*2, out_features]，并将其注册为模型参数。
            self.weights = Parameter(torch.FloatTensor(K, in_features * 2, out_features))
        elif backbone_type == 'gat':
            self.leakyrelu = nn.LeakyReLU()
            # 初始化一个权重矩阵 self.weights，其形状为 [K, in_features, out_features]，并将其注册为模型参数。
            self.weights = nn.Parameter(torch.zeros(K, in_features, out_features))
            # 初始化一个注意力权重矩阵 self.a，其形状为 [K, 2 * out_features, 1]，并将其注册为模型参数。
            self.a = nn.Parameter(torch.zeros(K, 2 * out_features, 1))
        self.K = K
        self.device = device
        self.variant = variant
        self.reset_parameters()

    def reset_parameters(self):
        # stdv 是用于初始化权重的标准差，计算方法为 1 / sqrt(out_features)。
        stdv = 1. / math.sqrt(self.out_features)
        # 使用 uniform_(-stdv, stdv) 方法将 self.weights 初始化为范围在 [-stdv, stdv] 之间的均匀分布。
        self.weights.data.uniform_(-stdv, stdv)
        # 如果 backbone_type 是 gat，使用 Xavier 均匀初始化方法初始化注意力参数 self.a，其中增益值为 1.414。
        if self.backbone_type == 'gat':
            nn.init.xavier_uniform_(self.a.data, gain=1.414)

    # 该方法用于执行特殊的稀疏矩阵乘法。具体来说，它将边索引和边权重转换为稀疏矩阵，然后与特征矩阵相乘。
    # adj：边索引，形状为 [2, E]，其中 E 是边的数量。
    # spm：边的权重，形状为 [E]。
    # size：稀疏矩阵的大小（例如 [N, N]），其中 N 是节点数。
    # h：输入特征矩阵，形状为 [N, D]，其中 D 是特征维度。
    def specialspmm(self, adj, spm, size, h):
        # 使用 torch_sparse 库中的 SparseTensor 类创建一个稀疏矩阵 adj。
        # row=adj[0] 和 col=adj[1] 提取边的起始节点和结束节点索引。
        # value=spm 使用边的权重。
        # sparse_sizes=size 指定稀疏矩阵的大小。
        adj = SparseTensor(row=adj[0], col=adj[1], value=spm, sparse_sizes=size)
        return matmul(adj, h)

    # x：输入特征矩阵，形状为 [N, D]，其中 N 是节点数，D 是特征维度。
    # adj：边索引，形状为 [2, E]，其中 E 是边的数量。
    # e：环境编码矩阵，形状为 [N, K]，其中 K 是卷积层数。
    # weights：卷积层的权重矩阵，默认为 None。
    def forward(self, x, adj, e, weights=None):
        # 如果 weights 为 None，则使用类定义的权重矩阵 self.weights。
        if weights == None:
            weights = self.weights
        # 如果 backbone_type 是 gcn：
        # 如果 variant 为 False，调用 gcn_conv(x, adj) 计算图卷积，得到新的节点表示 hi。
        # 否则，将 adj 转换为稀疏张量，并执行稀疏矩阵乘法，得到新的节点表示 hi。
        if self.backbone_type == 'gcn':
            if not self.variant:
                hi = gcn_conv(x, adj)
            else:
                adj = torch.sparse_coo_tensor(adj, torch.ones(adj.shape[1]).to(self.device),
                                              size=(x.shape[0], x.shape[0])).to(self.device)
                hi = torch.sparse.mm(adj, x)
            # 将卷积结果 hi 与输入特征 x 连接，形成形状为 [N, 2D] 的张量。
            hi = torch.cat([hi, x], 1)
            # 将 hi 扩展为三维张量，形状为 [K, N, 2D]，重复 K 次。
            # unsqueeze(dim) 会在指定的 dim 位置增加一个维度。在这里，unsqueeze(0) 会在第 0 维增加一个新的维度。
            # 如果 hi 的原始形状是 [N, 2D]，那么 hi.unsqueeze(0) 的形状将变成 [1, N, 2D]。
            # repeat 方法会沿指定的维度复制张量。在这里，repeat(self.K, 1, 1) 表示沿第 0 维复制 self.K 次，沿第 1 维和第 2 维不复制（保持原样）。
            # 如果 hi 的形状是 [1, N, 2D]，那么 hi.repeat(self.K, 1, 1) 的形状将变成 [K, N, 2D]。
            hi = hi.unsqueeze(0).repeat(self.K, 1, 1)  # [K, N, D*2]
            # 使用权重 weights 进行矩阵乘法，得到形状为 [K, N, D] 的张量 outputs。
            # //todo [K,N,2D]*[K,in_features*2,out_features]  [K,N,2D]*[K,2D,D]  不确定是哪个
            outputs = torch.matmul(hi, weights)  # [K, N, D]
            print(f"hi shape: {hi.shape},weight shape: {weights.shape},outputs shape: {outputs.shape}")
            # 转置 outputs，将形状转换为 [N, K, D]。
            # torch.transpose(input, dim0, dim1) 方法用于交换张量的两个维度。它不改变数据本身，只改变数据的视图。
            outputs = outputs.transpose(1, 0)  # [N, K, D]
        elif self.backbone_type == 'gat':
            # 如果 backbone_type 是 gat：
            # 将 x 扩展为三维张量，形状为 [K, N, D]。
            # 使用权重 weights 进行矩阵乘法，得到形状为 [K, N, D] 的张量 h。
            # 获取节点数 N。
            # 移除自环，并添加自环。
            xi = x.unsqueeze(0).repeat(self.K, 1, 1)  # [K, N, D]
            # //todo [K,N,2D]*[K, in_features, out_features]  [K,N,D]*[K,D,D]  不确定是哪个
            h = torch.matmul(xi, weights)  # [K, N, D]
            N = x.size()[0]
            adj, _ = remove_self_loops(adj)
            adj, _ = add_self_loops(adj, num_nodes=N)
            # 将边的起始节点和结束节点的特征连接，形成形状为 [K, E, 2D] 的张量 edge_h。
            # adj 是一个形状为 [2, E] 的二维张量，其中 E 是边的数量。adj[0, :] 和 adj[1, :] 分别表示边的起始节点和结束节点的索引。
            # adj[0, :] 提取所有边的起始节点索引，形状为 [E]。
            # h[:, adj[0, :], :] 提取 h 中这些起始节点的特征，得到的张量形状为 [K, E, D]，即每条边的起始节点特征。
            # adj[1, :] 提取所有边的结束节点索引，形状为 [E]。
            # h[:, adj[1, :], :] 提取 h 中这些结束节点的特征，得到的张量形状为 [K, E, D]，即每条边的结束节点特征。
            # (h[:, adj[0, :], :], h[:, adj[1, :], :]) 表示将起始节点特征和结束节点特征沿第 2 维度（特征维度）进行连接。
            # 连接后的张量 edge_h 形状为 [K, E, 2*D]，其中每条边的特征由起始节点和结束节点的特征拼接而成。
            edge_h = torch.cat((h[:, adj[0, :], :], h[:, adj[1, :], :]), dim=2)  # [K, E, 2*D]
            # 计算注意力系数：
            # 将 edge_h 与注意力参数 self.a 相乘，并使用 LeakyReLU 激活函数。
            # 移除第三个维度，得到形状为 [K, E] 的张量 logits。
            # edge_h 的形状为 [K, E, 2*D]。
            # self.a 的形状为 [K, 2*D, 1]。[K, 2 * out_features, 1]
            # 对于每个 k，edge_h[k] 的形状为 [E, 2*D]，self.a[k] 的形状为 [2*D, 1]。矩阵乘法 edge_h[k] @ self.a[k] 的结果是一个形状为 [E, 1] 的张量。
            # 由于 K 个这样的矩阵乘法是并行进行的，最终结果 logits 的形状为 [K, E, 1]。
            # squeeze 方法用于移除张量中大小为 1 的维度。logits 的形状为 [K, E, 1]，其中第 2 维度的大小为 1。
            # logits.squeeze(2) 将移除第 2 维度，得到的张量形状为 [K, E]。
            logits = self.leakyrelu(torch.matmul(edge_h, self.a)).squeeze(2)
            # 计算归一化的注意力权重：
            # 对 logits 进行最大值归一化，防止数值溢出。
            # 使用 exp 函数计算注意力权重。
            # logits 的形状为 [K, E]。
            # torch.max(logits, dim=1, keepdim=True) 在每一行中找到最大值。 keepdim：是否保持原始的维度。
            # logits_max 的形状为 [K, 1]，因为 keepdim=True 保持了第 1 维度的原始大小。
            logits_max, _ = torch.max(logits, dim=1, keepdim=True)
            # 在计算 softmax 时，通过减去最大值可以防止指数运算中的数值溢出。这是因为指数函数的值增长非常快，减去最大值可以使得计算在数值上更稳定。
            # logits_max 的形状为 [K, 1]，可以自动广播（broadcast）到 [K, E] 的形状，与 logits 进行元素级减法。
            # logits - logits_max 计算每个 logits 元素减去对应行的最大值，结果形状为 [K, E]。
            # torch.exp(logits - logits_max) 计算指数，结果形状为 [K, E]。
            # 数值稳定性：通过减去每行的最大值，可以防止指数运算中的数值溢出，提高计算的稳定性。
            # 归一化：在后续的 softmax 计算中，需要将 logits 归一化为概率分布。这一步是 softmax 计算的前奏。
            edge_e = torch.exp(logits - logits_max)  # [K, E]

            # 初始化一个空列表 outputs。
            outputs = []
            # 设置一个小的数值 eps 防止除零错误。
            eps = 1e-8
            # k：域的数量，每个域对应一个图卷积滤波器
            for k in range(self.K):
                # 提取第 k 层的注意力权重 edge_e_k。
                edge_e_k = edge_e[k, :]  # [E]
                # 使用 specialspmm 计算每个节点的归一化注意力权重和特征聚合结果。
                # specialspmm 方法计算稀疏邻接矩阵与全1张量的乘法，得到每个节点的注意力权重之和，这在归一化注意力机制中非常重要。
                # 防止数值错误：通过添加一个小的 eps 值，可以防止后续计算中的除零错误。
                e_expsum_k = self.specialspmm(adj, edge_e_k, torch.Size([N, N]), torch.ones(N, 1).cuda()) + eps
                # assert not torch.isnan(e_expsum_k).any() 确保 e_expsum_k 中没有任何 NaN 值，如果有则抛出断言错误。
                assert not torch.isnan(e_expsum_k).any()

                # 注意力权重的应用：通过稀疏矩阵乘法，可以将注意力权重应用到特征聚合过程中。
                # 归一化：逐元素除法实现了特征的归一化处理，使得每个节点的特征表示在数值上更加稳定，避免了注意力权重的累积导致的数值爆炸。
                hi_k = self.specialspmm(adj, edge_e_k, torch.Size([N, N]), h[k])
                # torch.div(input, other) 用于执行逐元素除法。
                # hi_k 的形状为 [N, D]。
                # e_expsum_k 的形状为 [N, 1]，可以自动广播（broadcast）到 [N, D] 的形状。
                # torch.div(hi_k, e_expsum_k) 对 hi_k 的每个元素进行除法操作，除以对应行的 e_expsum_k 值。
                hi_k = torch.div(hi_k, e_expsum_k)  # [N, D]
                # 将结果添加到 outputs 列表中。
                outputs.append(hi_k)
            # 将 outputs 列表堆叠为一个形状为 [N, K, D] 的张量。
            # outputs 是一个包含 K 个形状为 [N, D] 的张量的列表。
            # torch.stack(outputs, dim=1) 在第 1 维度插入一个新的维度，将这些张量堆叠起来。
            # 堆叠后的张量形状为 [N, K, D]，其中 K 是列表中张量的数量。
            outputs = torch.stack(outputs, dim=1)  # [N, K, D]

        # 将环境编码矩阵 e 扩展为三维张量，并重复 self.out_features 次，形成形状为 [N, K, D] 的张量 es。
        # e 是一个形状为 [N, K] 的二维张量，其中 N 是节点数，K 是卷积层数或头数。
        # self.out_features 表示输出特征的维度，这里等于 D。
        # e.unsqueeze(2) 在第 2 维度（索引为 2）增加一个新的维度，使得 e 的形状从 [N, K] 变为 [N, K, 1]。
        # .repeat(1, 1, self.out_features) 将张量在第 2 维度重复 self.out_features 次，使得形状变为 [N, K, self.out_features]。
        es = e.unsqueeze(2).repeat(1, 1, self.out_features)  # [N, K, D]
        # 将 es 与 outputs 逐元素相乘，并在第一个维度上求和，得到形状为 [N, D] 的最终输出 output。
        # torch.mul(es, outputs) 对 es 和 outputs 进行逐元素相乘，得到一个形状为 [N, K, D] 的张量。
        # torch.sum(torch.mul(es, outputs), dim=1) 沿第 1 维度（索引为 1）进行求和，得到一个形状为 [N, D] 的张量。
        # 最终实现了特征的加权聚合和调整，使得不同卷积层或头的输出特征被整合到一起，得到每个节点的最终特征表示。
        output = torch.sum(torch.mul(es, outputs), dim=1)  # [N, D]

        # 如果使用残差连接，将输入特征 x 加到输出 output 上。
        if self.residual:
            output = output + x

        return output


class CaNet(nn.Module):
    # d：输入特征的维度。
    # c：输出类别的数量。
    # args：包含各种参数的对象。
    # device：指定计算设备（如 GPU 或 CPU）。
    def __init__(self, d, c, args, device):
        super(CaNet, self).__init__()
        # self.convs 是一个 ModuleList，用于存储多个卷积层。
        # 遍历 range(args.num_layers)，根据指定的层数添加多个 CaNetConv 层。
        # CaNetConv 的参数：
        # args.hidden_channels：隐藏层特征的维度。
        # args.K：域的数量，每个域对应一个图卷积滤波器
        # args.backbone_type：选择使用的卷积类型（如 GCN 或 GAT）。
        # residual=True：使用残差连接。
        # device：计算设备。
        # variant=args.variant：使用变体版本的卷积。
        self.convs = nn.ModuleList()
        for _ in range(args.num_layers):
            self.convs.append(
                CaNetConv(args.hidden_channels, args.hidden_channels, args.K, backbone_type=args.backbone_type,
                          residual=True, device=device, variant=args.variant))
        # self.fcs 是一个 ModuleList，用于存储多个全连接层。
        # 第一个全连接层将输入特征从维度 d 转换为隐藏层特征维度 args.hidden_channels。
        # 第二个全连接层将隐藏层特征维度转换为输出类别数量 c。
        self.fcs = nn.ModuleList()
        self.fcs.append(nn.Linear(d, args.hidden_channels))
        self.fcs.append(nn.Linear(args.hidden_channels, c))
        # self.env_enc 是一个 ModuleList，用于存储多个环境编码层。
        self.env_enc = nn.ModuleList()
        # 遍历 range(args.num_layers)，根据指定的层数添加多个环境编码层。
        for _ in range(args.num_layers):
            # 根据 args.env_type 决定使用哪种类型的环境编码层：
            if args.env_type == 'node':
                # 如果是 node，使用 Linear 层将隐藏层特征维度转换为卷积操作的重复次数 args.K。
                # 节点级别的环境编码：当 env_type 为 node 时，表示我们只需要对每个节点的特征进行简单的线性变换，而不需要考虑节点之间的连边关系。
                # 适用场景：这种情况适用于那些节点特征独立于图结构的情形。通过线性层 nn.Linear，我们可以对每个节点的特征进行线性变换，以得到环境编码。
                # 实现：nn.Linear(args.hidden_channels, args.K)，将节点特征从维度 args.hidden_channels 转换到维度 args.K。
                self.env_enc.append(nn.Linear(args.hidden_channels, args.K))
            elif args.env_type == 'graph':
                # 如果是 graph，使用 GraphConvolutionBase 层将隐藏层特征维度转换为卷积操作的重复次数 args.K。
                # 图级别的环境编码：当 env_type 为 graph 时，表示我们需要考虑节点之间的连边关系，对整个图进行卷积操作。
                # 适用场景：这种情况适用于节点特征受图结构影响的情形。通过图卷积层 GraphConvolutionBase，我们可以在考虑节点连边关系的基础上对节点特征进行变换。
                # 实现：GraphConvolutionBase(args.hidden_channels, args.K, residual=True)，将节点特征从维度 args.hidden_channels
                # 转换到维度 args.K，并考虑图结构信息。
                self.env_enc.append(GraphConvolutionBase(args.hidden_channels, args.K, residual=True))
            else:
                # 如果不是上述类型，抛出未实现错误 NotImplementedError。
                raise NotImplementedError
        # self.act_fn：定义激活函数为 ReLU。
        self.act_fn = nn.ReLU()
        # self.dropout：设置 dropout 概率。
        self.dropout = args.dropout
        # self.num_layers：设置卷积层数。
        self.num_layers = args.num_layers
        # self.tau：设置 Gumbel-Softmax 温度参数。
        self.tau = args.tau
        # self.env_type：设置环境类型。
        self.env_type = args.env_type
        # self.device：设置计算设备。
        self.device = device

    # reset_parameters 方法用于初始化模型各层的参数。
    def reset_parameters(self):
        # 遍历 self.convs 列表中的每个卷积层 conv。
        # 调用每个卷积层 conv 的 reset_parameters 方法，以初始化其权重和偏置参数。
        for conv in self.convs:
            conv.reset_parameters()
        # 遍历 self.fcs 列表中的每个全连接层 fc。
        # 调用每个全连接层 fc 的 reset_parameters 方法，以初始化其权重和偏置参数。
        for fc in self.fcs:
            fc.reset_parameters()
        # 遍历 self.env_enc 列表中的每个环境编码层 enc。
        # 调用每个环境编码层 enc 的 reset_parameters 方法，以初始化其权重和偏置参数。
        for enc in self.env_enc:
            enc.reset_parameters()

    # x：输入特征矩阵，形状为 [N, D]，其中 N 是节点数，D 是输入特征维度。
    # adj：边的索引矩阵，形状为 [2, E]，其中 E 是边的数量。
    # idx：索引，用于选择特定节点进行操作，默认为 None。
    # training：布尔值，表示是否在训练模式下，默认为 False。
    def forward(self, x, adj, idx=None, training=False):
        self.training = training
        # 使用 F.dropout 对输入特征 x 进行 Dropout 处理，Dropout 概率为 self.dropout，仅在训练模式下生效。
        x = F.dropout(x, self.dropout, training=self.training)
        # 使用第一个全连接层 self.fcs[0] 处理 Dropout 后的输入特征 x，并通过激活函数 self.act_fn，得到中间表示 h。
        # [N,D] 变为[N,hidden_channels]
        h = self.act_fn(self.fcs[0](x))
        # 克隆中间表示 h，得到 h0，用于后续的残差连接。
        h0 = h.clone()

        # 初始化正则化损失 reg 为 0。
        reg = 0
        # 遍历每个卷积层 con 和对应的索引 i，对 self.convs 列表中的每个卷积层进行前向传播。
        for i, con in enumerate(self.convs):
            # 使用 F.dropout 对中间表示 h 进行 Dropout 处理，Dropout 概率为 self.dropout，仅在训练模式下生效。
            h = F.dropout(h, self.dropout, training=self.training)
            # 如果在训练模式下：
            if self.training:
                # 如果 env_type 是 node，则使用当前环境编码层 self.env_enc[i] 对 h 进行线性变换，得到 logit。
                # 否则，使用当前环境编码层 self.env_enc[i] 对 h 和 adj 进行图卷积变换，得到 logit。[N,K]
                if self.env_type == 'node':
                    logit = self.env_enc[i](h)
                else:
                    logit = self.env_enc[i](h, adj, h0)
                # 使用 Gumbel-Softmax 采样 logit，温度参数为 self.tau，得到环境编码 e。
                e = F.gumbel_softmax(logit, tau=self.tau, dim=-1)
                # 计算正则化损失，并累加到 reg 中。
                reg += self.reg_loss(e, logit)
            # 如果不在训练模式下：
            else:
                # 如果 env_type 是 node，则对 h 使用 Softmax，得到环境编码 e。
                # 否则，对 h 和 adj 使用 Softmax，得到环境编码 e。
                if self.env_type == 'node':
                    e = F.softmax(self.env_enc[i](h), dim=-1)
                else:
                    e = F.softmax(self.env_enc[i](h, adj, h0), dim=-1)
            # 使用当前卷积层 con 对 h、adj 和环境编码 e 进行卷积操作，并通过激活函数 self.act_fn，更新中间表示 h
            h = self.act_fn(con(h, adj, e))

        # 使用 F.dropout 对最终的中间表示 h 进行 Dropout 处理，Dropout 概率为 self.dropout，仅在训练模式下生效。
        h = F.dropout(h, self.dropout, training=self.training)
        # 使用最后一个全连接层 self.fcs[-1] 对 Dropout 后的中间表示 h 进行线性变换，得到最终输出 out。
        # out的维度是[节点数,类别]
        out = self.fcs[-1](h)
        # 如果在训练模式下：返回输出 out 和正则化损失 reg 除以层数 self.num_layers 的平均值。
        # 如果不在训练模式下：仅返回输出 out。
        if self.training:
            return out, reg / self.num_layers
        else:
            return out

    # 这个方法用于计算正则化损失
    # 正则化：通过计算 Gumbel-Softmax 或 Softmax 编码 z 与 log-softmax log_pi 的加权对数概率，可以对模型的环境编码进行正则化，防止过拟合。
    # 信息熵：这种计算方式类似于信息熵，可以衡量环境编码的分布情况，鼓励模型在不同环境之间进行适当的权衡和选择。
    # 数值稳定性：使用 log-sum-exp 技术计算 log-softmax，可以提高数值稳定性，防止计算过程中的溢出或下溢。
    def reg_loss(self, z, logit, logit_0=None):
        # logit：输入的 logits 张量，形状为 [N, K]，其中 N 是节点数，K 是类别数或环境数量。
        # torch.logsumexp(logit, dim=-1, keepdim=True)：计算 logits 在最后一维上的 log-sum-exp，结果形状为 [N, 1]。
        # logsumexp 计算公式：logsumexp(x) = log(sum(exp(x)))，用于数值稳定性。
        # dim=-1 表示在最后一维上进行操作。
        # keepdim=True 表示保持原始维度。
        # .repeat(1, logit.size(1))：将 [N, 1] 形状的张量重复为 [N, K]，以匹配 logits 的形状。
        # logit - torch.logsumexp(logit, dim=-1, keepdim=True).repeat(1, logit.size(1))：计算 log-softmax，得到对数概率 log_pi，形状为 [N, K]。
        log_pi = logit - torch.logsumexp(logit, dim=-1, keepdim=True).repeat(1, logit.size(1))
        # z：Gumbel-Softmax 后的环境编码，形状为 [N, K]。
        # torch.mul(z, log_pi)：逐元素相乘，计算每个节点和每个类别的加权对数概率，形状为 [N, K]。
        # torch.sum(torch.mul(z, log_pi), dim=1)：在最后一维上求和，得到每个节点的总加权对数概率，形状为 [N]。
        # torch.mean(torch.sum(torch.mul(z, log_pi), dim=1))：对所有节点的总加权对数概率取平均，得到正则化损失的标量值。
        return torch.mean(torch.sum(
            torch.mul(z, log_pi), dim=1))

    # 这个方法用于计算监督损失
    def sup_loss_calc(self, y, pred, criterion, args):
        # 检查当前数据集是否为 'twitch' 或 'elliptic'。
        # 这些数据集可能是多标签分类任务，需要特别处理标签和损失计算。
        if args.dataset in ('twitch', 'elliptic'):
            # 如果 y 的形状为 [N, 1]（即每个样本只有一个标签），需要将标签转换为 one-hot 编码。
            if y.shape[1] == 1:
                # F.one_hot(y, y.max() + 1)：将标签 y 转换为 one-hot 编码。
                # y：输入标签张量，形状为 [N, 1]，其中 N 是样本数量。
                # y.max() + 1：类别的数量。y.max() 返回标签张量中的最大值，加1得到类别数量。

                # 假设 y 是一个形状为 [N, 1] 的张量，如下所示：
                # y = torch.tensor([[0], [1], [2], [1]])  # 形状为 [N, 1]
                # 那么 y.max() + 1 为 3，表示有 3 个类别。
                # F.one_hot(y, y.max() + 1) 将 y 转换为 one-hot 编码：
                # one_hot_y = F.one_hot(y, num_classes=3)
                # 结果 one_hot_y 形状为 [N, 1, 3]，其中 3 是类别数量。

                # squeeze(1)：移除第二个维度（即 [N, 1] 变为 [N]）。squeeze(1)代表若第二维度值为1则去除第二维度
                # 最终true_label变为[N, num_classes]
                true_label = F.one_hot(y, y.max() + 1).squeeze(1)
            # 如果 y 已经是多标签格式（如 [N, num_classes]），直接使用 y 作为 true_label。
            else:
                true_label = y
            # 使用给定的损失函数 criterion 计算预测值 pred 和真实标签 true_label 之间的损失。
            # true_label.squeeze(1).to(torch.float)：将真实标签移除第二个维度并转换为浮点数。
            # 这样做是为了确保预测值和真实标签的格式匹配，适用于二分类或多标签分类任务。
            loss = criterion(pred, true_label.squeeze(1).to(torch.float))
        # 对于其他数据集，假设是多分类任务，直接使用 log_softmax 计算预测值的对数概率。
        else:
            # F.log_softmax(pred, dim=1)：对预测值 pred 在类别维度上计算 log-softmax。
            # log-softmax 是一种将输入张量中的每个元素转换为对数概率的方法，常用于分类任务中。
            out = F.log_softmax(pred, dim=1)
            # 将标签 y 移除第二个维度。
            # squeeze(1)代表若第二维度值为1则去除第二维度
            target = y.squeeze(1)
            # 使用给定的损失函数 criterion 计算 log_softmax 预测值 out 和真实标签 target 之间的损失。
            loss = criterion(out, target)
        return loss

    # 这个方法用于计算总损失，包括监督学习损失和正则化损失
    def loss_compute(self, d, criterion, args):
        # 调用模型的 forward 方法进行前向传播，计算预测的 logits 和正则化损失。
        # d.x：输入特征矩阵，形状为 [N, D]，其中 N 是节点数，D 是特征维度。
        # 参数：
        # d.edge_index：图的边索引矩阵，表示图结构。
        # idx=d.train_idx：训练节点的索引，用于选择训练节点进行计算。
        # training=True：指示模型处于训练模式，影响 Dropout 和 BatchNorm 等层的行为。
        # 返回值：
        # logits：模型的输出，形状为 [N, C]，其中 C 是类别数。
        # reg_loss：正则化损失，用于防止过拟合。
        logits, reg_loss = self.forward(d.x, d.edge_index, idx=d.train_idx, training=True)
        # 调用 sup_loss_calc 方法计算监督学习损失。
        # 参数：
        # d.y[d.train_idx]：训练节点的真实标签，形状为 [M] 或 [M, 1]，其中 M 是训练节点数。
        # logits[d.train_idx]：训练节点的预测 logits，形状为 [M, C]。
        # criterion：损失函数，用于计算预测值和真实标签之间的损失。
        # args：包含各种参数的对象。
        # 返回值：
        # sup_loss：监督学习损失，表示模型预测值与真实标签之间的差异。
        sup_loss = self.sup_loss_calc(d.y[d.train_idx], logits[d.train_idx], criterion, args)
        # 计算总损失，包含监督学习损失和正则化损失的加权和。
        # args.lamda：正则化损失的权重系数，用于平衡监督学习损失和正则化损失的影响。
        # 返回值：
        # loss：总损失，用于优化模型参数。
        loss = sup_loss + args.lamda * reg_loss
        return loss
